import numpy as np
import time
import logging

from gym import spaces, core
# core.Env是gym的环境基类,自定义的环境就是根据自己的需要重写其中的方法；
# 必须要重写的方法有:
# __init__()：构造函数
# reset()：初始化环境
# step()：环境动作,即环境对agent的反馈
# render()：如果要进行可视化则实现
import sys
sys.path.append('/apps/users/dl02/George/zeromq/rosproxy_python')
import rosproxy



MAX_V = 1.
MAX_W = 1.

class FarmTractor(core.Env):
    def __init__(self):
        self.action_space = spaces.Box(low=np.array([-MAX_V, -MAX_W]), high=np.array([MAX_V, MAX_W]))  # 动作空间
        self.calibration_time = rosproxy.message.std_msgs.Time()
        self.truth_heading = None
        self.truth_position = None
        self.truth_v = None
        self.truth_w = None
        self.time_step = None
        # self.observation_space = spaces.Box(low=-1, high=1, shape=(1,))  # 状态空间

        # 必须要的全局初始化函数
        rosproxy.proxy.Setup('tianjin', 'tcp://172.22.60.113:6658', 'tcp://172.22.60.113:6587')
        # 创建一个topic
        subcribe = rosproxy.proxy.Topic()
        # 初始化一个topic
        subcribe.Init('Tractor01', '/tianjin/Tractor01_00/data', 'custom_msgs/Data', 1)
        # 订阅消息，有消息到的到就会调用回调函数
        subcribe.Subscribe(self.callback)

        # creat a topic for controlling
        self.ControlPublish = rosproxy.proxy.Topic()
        self.ControlPublish.Init('Tractor01', '/tianjin/Tractor01_00/cmd_vel', 'geometry_msgs/TwistStamped', 1)
        self.ControlPublish.Advertise()

        # creat a topic for controlling
        self.ResetPublish = rosproxy.rosproxy.Topic()
        self.ResetPublish.Init('Tractor01', '/tianjin/reset', 'std_msgs/Time', 1)
        self.ResetPublish.Advertise()

    def callback(self, aData):
        # print('仿真世界时间: ', aData.global_time._Sec, '.', aData.global_time._NSec)
        # print('控制校对时间: ', aData.control_time._Sec, '.', aData.control_time._NSec)
        self.truth_heading = (aData.pose.orientation.x, aData.pose.orientation.y,aData.pose.orientation.z, aData.pose.orientation.w)
        self.truth_position = (aData.pose.orientation.x, aData.pose.orientation.y,aData.pose.orientation.z, aData.pose.orientation.w)
        self.truth_v = (aData.twist.linear.x, aData.twist.linear.y, aData.twist.linear.z)
        self.truth_w = (aData.twist.angular.x, aData.twist.angular.y, aData.twist.angular.z)
        self.calibration_time.data = aData.control_time

    def reset(self):
        # 新建一个消息体
        aReset = rosproxy.message.std_msgs.Time()

        # print(dir(aReset)) # 查看当前消息结构体组成
        # aReset.data._NSec = 0  # 重置时间：纳秒 %仿真重置完成返回
        aReset.data._Sec = 0  # 重置时间：秒 %仿真重置完成返回

        while self.calibration_time.data._Sec != 0:
            self.ResetPublish.Publish(aReset)  # 发布指令
            logging.info('重置仿真')
            # print(calibration_time.data._NSec) # 等待校对时间吻合
            time.sleep(1)
        logging.info('重置成功')
        ob = [self.truth_heading, self.truth_position, self.truth_v, self.truth_w]
        return ob

    def control(self, LV, AV, time_step):
        # 新建一个消息体
        aCmd_vel = rosproxy.message.geometry_msgs.TwistStamped()

        # print(dir(aCmd_vel)) # 查看当前消息结构体组成
        aCmd_vel.twist.linear.x = LV  # 控制输入：线速度
        aCmd_vel.twist.angular.z = AV  # 控制输入：角速度
        # aCmd_vel.header.time._NSec = 123456  # 校对时间：纳秒 %下次仿真冻结返回
        aCmd_vel.header.time._Sec = time_step # 校对时间：秒 %下次仿真冻结返回

        self.ControlPublish.Publish(aCmd_vel)  # 发布指令

    def step(self, action):
        while self.calibration_time.data != self.time_step + 1:
            self.control(1, 0, self.time_step)
            time.sleep(0.25)  # 休眠若干毫秒
        self.time_step += 1
        reward = self._get_reward()
        done = self._get_done()
        obs = [self.truth_heading, self.truth_position, self.truth_v, self.truth_w]
        info = {}  # 用于记录训练过程中的环境信息,便于观察训练状态
        return obs, reward, done, info

    # 根据需要设计相关辅助函数
    # def _get_observation(self, action):
    #     ...
    #     return obs

    # TODO
    def _get_reward(self):
        reward = 0
        return reward

    # TODO
    def _get_done(self):
        done = False
        return done
